package wge

import (
	"net"
	"net/http"
	"sync"
	"time"
	"log"
)

type Server interface {
	Serve(string, http.Handler) error
	Close() error
}

type server struct {
	l		net.Listener
	srv 	*http.Server
	quit	chan int
	mu		sync.Mutex
	conns 	map[string]net.Conn
	wg 		sync.WaitGroup
}

func (s *server)Serve(host string, h http.Handler) error {
	s.srv = &http.Server{Addr: host, Handler: h}
	addr := s.srv.Addr
	if addr == "" {
		addr = ":http"
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		logger.Println(err)
	}
	s.l = l
	s.quit = make(chan int)
	s.conns = make(map[string]net.Conn)

	// 处理客户端连接状态，主要是针对 keep
	s.srv.ConnState = func(conn net.Conn, state http.ConnState) {
		switch state {
		case http.StateNew:
			s.wg.Add(1)
		case http.StateActive:
			s.mu.Lock()
			delete(s.conns, conn.RemoteAddr().String())
			s.mu.Unlock()
		case http.StateIdle:
			select {
			case <- s.quit:		// 如果关闭了连接，直接关掉客户端连接，否则加入连接池
				conn.Close()
			default:
				s.mu.Lock()
				s.conns[conn.RemoteAddr().String()] = conn
				s.mu.Unlock()
			}
		case http.StateHijacked, http.StateClosed:
			s.mu.Lock()
			delete(s.conns, conn.RemoteAddr().String())
			s.mu.Unlock()
			s.wg.Done()
		}
	}
	return s.srv.Serve(s.l)
	//err := http.ListenAndServe(host, g)
}

func (s *server)Close() error {
	mu.Lock()
	if !isStarted {
		log.Println("server isn't started!")
		mu.Unlock()
		return nil
	}
	mu.Unlock()
	// 通过关闭 channel 来通知客户端不要进入kepp状态
	close(s.quit)

	// 关闭keepalived，请求返回的时候会带上Close header。客户端就知道要close掉connection了。
	s.srv.SetKeepAlivesEnabled(false)

	s.mu.Lock()
	if err := s.l.Close(); err != nil {
		logger.Println(err)
		return err
	}
	// 将当前 idle 的 conns 设置 read timeout，以便后续关闭
	t := time.Now().Add(100 * time.Millisecond)
	for _, c := range s.conns {
		c.SetReadDeadline(t)
	}
	s.mu.Unlock()

	// 等待所有连接关闭 sync WaitGroup
	s.wg.Wait()

	return nil
}
